# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
from sklearn.preprocessing import LabelEncoder
from lightgbm import LGBMClassifier
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score
import warnings
from tqdm import tqdm

warnings.filterwarnings('ignore')

# 读取数据
train_data = pd.read_csv(r'dataTrain.csv')
test_data = pd.read_csv(r'dataB.csv')
submission = pd.read_csv(r'submit_example_B.csv')

print(f'训练数据形状 = {train_data.shape}\n测试数据形状  = {test_data.shape}')

# 构造特征
train_data['f47'] = train_data['f1'] * 10 + train_data['f2']
test_data['f47'] = test_data['f1'] * 10 + test_data['f2']

# 暴力特征组合
loc_f = ['f1', 'f2', 'f4', 'f5', 'f6']
for df in [train_data, test_data]:
    for i in range(len(loc_f)):
        for j in range(i + 1, len(loc_f)):
            df[f'{loc_f[i]}+{loc_f[j]}'] = df[loc_f[i]] + df[loc_f[j]]
            df[f'{loc_f[i]}-{loc_f[j]}'] = df[loc_f[i]] - df[loc_f[j]]
            df[f'{loc_f[i]}*{loc_f[j]}'] = df[loc_f[i]] * df[loc_f[j]]
            df[f'{loc_f[i]}/{loc_f[j]}'] = df[loc_f[i]] / (df[loc_f[j]] + 1)

# 离散化
all_f = [f'f{idx}' for idx in range(1, 47) if idx != 3]
for df in [train_data, test_data]:
    for col in all_f:
        df[f'{col}_log'] = df[col].apply(lambda x: int(np.log(x)) if x > 0 else 0)

# 特征交叉
log_f = [f'f{idx}_log' for idx in range(1, 47) if idx != 3]
for df in [train_data, test_data]:
    for i in range(len(log_f)):
        for j in range(i + 1, len(log_f)):
            df[f'{log_f[i]}_{log_f[j]}'] = df[log_f[i]] * 10000 + df[log_f[j]]

cat_columns = ['f3']
data = pd.concat([train_data, test_data])

for col in cat_columns:
    lb = LabelEncoder()
    data[col] = lb.fit_transform(data[col])

# 准备训练数据
feature_columns = [col for col in train_data.columns if col not in ['id', 'label', 'f3']]
target = 'label'
X_train = train_data[feature_columns]
y_train = train_data[target]
test = test_data[feature_columns]

# 模型训练函数
def model_train(model, model_name, kfold=10):
    oof_preds = np.zeros((X_train.shape[0]))
    test_preds = np.zeros(test.shape[0])
    skf = StratifiedKFold(n_splits=kfold)

    for k, (train_index, val_index) in enumerate(skf.split(X_train, y_train)):
        x_train, x_val = X_train.iloc[train_index], X_train.iloc[val_index]
        y_train_fold, y_val = y_train.iloc[train_index], y_train.iloc[val_index]

        model.fit(x_train, y_train_fold)

        y_pred = model.predict_proba(x_val)[:, 1]
        oof_preds[val_index] = y_pred.ravel()
        auc = roc_auc_score(y_val, y_pred)
        print("模型 = %s, KFold = %d, 验证 AUC = %.4f" % (model_name, k, auc))
        test_fold_preds = model.predict_proba(test)[:, 1]
        test_preds += test_fold_preds.ravel()

    print("整体模型 = %s, AUC = %.4f" % (model_name, roc_auc_score(y_train, oof_preds)))
    return test_preds / kfold

# 创建模型并训练
gbm = LGBMClassifier(
    objective='binary',
    boosting_type='gbdt',
    num_leaves=31,
    max_depth=-1,   # 不限制深度
    min_child_samples=20,  # 每个叶子节点的最少样本数
    colsample_bytree=0.8,
    subsample=0.8,  # 设定子样本比率
    learning_rate=0.1,  # 调整学习率
    n_estimators=100,
    metric='auc'
)

# 训练模型
preds = model_train(gbm, "LGBMClassifier", 10)

# 保存预测结果
submission['label'] = preds
submission.to_csv('submission2.csv', index=False)
